import os
from aishangba_org import AiShangBa_Org
from shuquge_com import Shuquge_Com


def aishangba_org():
    url = input('欢迎使用“新笔趣阁”小说下载器，请输入小说网址：\n')
    # url = 'https://www.aishangba.org/114_114140/'
    path = f'{os.path.abspath(".")}\\www.aishangba.org'
    book = AiShangBa_Org(url)
    book.init(path)
    print(book.__format__())


def shuquge_com():
    url = input('欢迎使用“书趣阁”小说下载器，请输入小说网址：\n')
    path = f'{os.path.abspath(".")}\\www.shuquge.com'
    book = Shuquge_Com(url)
    book.init(path)
    print(book.__format__())


def main():
    while True:
        print(
            """=================================
请输入要爬取的小说网站：
1.新笔趣阁\t2.书趣阁
0.退出程序
=================================""")
        num = int(input('请输入要爬取的小说网站编号：'))
        if num == 0:
            print('退出程序')
            break
        elif num == 1:
            aishangba_org()
        elif num == 2:
            shuquge_com()


if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print('\n程序异常退出: ', e)
